package function

import (
	"context"
	"github.com/Jeffail/tunny"
	"sync"
)

//NewWaitGroup maxConcurrent最大并发，count总数
func NewWaitGroup(maxConcurrent, count int) *waitGroup {
	p := tunny.NewFunc(maxConcurrent, func(f interface{}) interface{} {
		fc := f.(func() error)
		return fc()

	})
	ctx, cancel := context.WithCancel(context.Background())
	wg := &waitGroup{pool: p, ctx: ctx, cancel: cancel}
	if count > 0 {
		wg.wg.Add(count)
	}
	return wg
}

type waitGroup struct {
	pool    *tunny.Pool
	wg      sync.WaitGroup
	errOnce sync.Once
	err     error
	ctx     context.Context
	cancel  context.CancelFunc
}

func (wg *waitGroup) Run(f func() error) {
	go func() {
		defer wg.wg.Done()
		err, _ := wg.pool.ProcessCtx(wg.ctx, f)
		if err != nil {
			wg.errOnce.Do(func() {
				wg.err = err.(error)
				wg.cancel()
			})
		}
	}()
	return
}

func (wg *waitGroup) OneDone() {
	wg.wg.Done()
}

func (wg *waitGroup) Wait() error {
	wg.wg.Wait()
	wg.pool.Close()
	return wg.err
}
